#include <iostream>
#include <chrono>
#include <cmath>
#include <hip/hip_runtime.h>
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

#define K 256
#define pad 1 // (K-1)/2

__global__ void convolution_kernel(float *input, float *kernel, float *output, int width, int height)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;

    if (i >= height || j >= width)
        return;

    float sum = 0.0f;
    for (int kx = 0; kx < K; ++kx)
    {
        for (int ky = 0; ky < K; ++ky)
        {
            int x = i - pad + kx;
            int y = j - pad + ky;
            if (x >= 0 && x < height && y >= 0 && y < width)
            {
                sum += input[x * width + y] * kernel[kx * K + ky];
            }
        }
    }
    output[i * width + j] = sum;
}

void convolution_cpu(float *input, float *kernel, float *output, int width, int height)
{
    for (int i = 0; i < height; ++i)
    {
        for (int j = 0; j < width; ++j)
        {
            float sum = 0.0f;
            for (int kx = 0; kx < K; ++kx)
            {
                for (int ky = 0; ky < K; ++ky)
                {
                    int x = i - pad + kx;
                    int y = j - pad + ky;
                    if (x >= 0 && x < height && y >= 0 && y < width)
                    {
                        sum += input[x * width + y] * kernel[kx * K + ky];
                    }
                }
            }
            output[i * width + j] = sum;
        }
    }
}

int main()
{
    // 读取图像
    int width, height, channels;
    unsigned char *img_data = stbi_load("input.png", &width, &height, &channels, 0);
    if (!img_data) {
        std::cerr << "Error loading image" << std::endl;
        return -1;
    }

    size_t img_size = width * height;
    float *h_input = new float[img_size];
    
    // 转换为灰度图并归一化到0-1
    for (int i = 0; i < img_size; ++i) {
        unsigned char r = img_data[i * channels];
        unsigned char g = img_data[i * channels + 1];
        unsigned char b = img_data[i * channels + 2];
        h_input[i] = (0.299f*r + 0.587f*g + 0.114f*b) / 255.0f;
    }
    stbi_image_free(img_data);

    float h_kernel[K*K];
    for (int i = 0; i < K*K; ++i)
        h_kernel[i] = 1.0f / (K*K);

    float *h_output_cpu = new float[img_size];
    float *h_output_gpu = new float[img_size];

    // GPU计算
    float *d_input, *d_kernel, *d_output;
    hipMalloc(&d_input, img_size * sizeof(float));
    hipMalloc(&d_kernel, K*K * sizeof(float));
    hipMalloc(&d_output, img_size * sizeof(float));

    hipMemcpy(d_input, h_input, img_size * sizeof(float), hipMemcpyHostToDevice);
    hipMemcpy(d_kernel, h_kernel, K*K * sizeof(float), hipMemcpyHostToDevice);

    dim3 block(16, 16);
    dim3 grid((width + block.x - 1)/block.x, (height + block.y - 1)/block.y);

    hipEvent_t start, stop;
    hipEventCreate(&start);
    hipEventCreate(&stop);

    hipEventRecord(start);
    convolution_kernel<<<grid, block>>>(d_input, d_kernel, d_output, width, height);
    hipEventRecord(stop);
    hipEventSynchronize(stop);

    float gpu_time;
    hipEventElapsedTime(&gpu_time, start, stop);
    hipMemcpy(h_output_gpu, d_output, img_size * sizeof(float), hipMemcpyDeviceToHost);

    // CPU计算
    auto cpu_start = std::chrono::high_resolution_clock::now();
    convolution_cpu(h_input, h_kernel, h_output_cpu, width, height);
    auto cpu_end = std::chrono::high_resolution_clock::now();
    float cpu_time = std::chrono::duration<float, std::milli>(cpu_end - cpu_start).count();

    // 验证结果
    bool match = true;
    const float epsilon = 1e-5f;
    for (int i = 0; i < img_size; ++i)
    {
        if (fabs(h_output_cpu[i] - h_output_gpu[i]) > epsilon)
        {
            match = false;
            break;
        }
    }

    std::cout << "CPU time: " << cpu_time << "ms\n";
    std::cout << "GPU time: " << gpu_time << "ms\n";
    std::cout << "Results match: " << (match ? "Yes" : "No") << "\n";

    // 保存结果到文件
    unsigned char* output_img = new unsigned char[img_size];
    for (int i = 0; i < img_size; ++i)
        output_img[i] = static_cast<unsigned char>(h_output_gpu[i] * 255.0f);
    
    stbi_write_png("output_gpu.png", width, height, 1, output_img, width);
    delete[] output_img;

    // 释放资源
    delete[] h_input;
    delete[] h_output_cpu;
    delete[] h_output_gpu;
    hipFree(d_input);
    hipFree(d_kernel);
    hipFree(d_output);

    return 0;
}